#pragma once


#include <mutex>
#include <map>
#include <optional>

#include "txdnn_internal.h"

#include "macros/errors.h"
#include "operators/names.h"
#include "operators/invoker_cache.h"
#include "operators/kernel_info.h"
#include "backend/kernel_interface/kernel_scheduler.h"


// Empty impl from txdnnContext to txdnnHandle_t
struct txdnnContext {};

namespace txdnn {

struct HandleImpl;

struct Handle : txdnnContext {
public:
  Handle();
  // Handle(txdnnAcceleratorQueue_t stream);
  Handle(Handle&& other) _TX_NOEXCEPT;
  virtual ~Handle();

  // stream 
  // txdnnAcceleratorQueue_t getStream() const {}
  // void setStream(txdnnAcceleratorQueue_t stream) const;

  // Allocator
  void setAllocator() const {}

  // Kernel
  txdnn::backend::KernelInvoke addKernel(const std::string& algorithm,
                                         const std::string& cfg,
                                         const fs::path& program_path,
                                         const std::string& kernel_name,
                                         const std::vector<size_t>& vld,
                                         const std::vector<size_t>& vgd,
                                         const std::string& params,
                                         std::size_t cache_idx          = 0,
                                         const std::string& kernel_src  = "") const;

  Invoker prepareInvoker(const InvokerFactory& factory,
                         const std::vector<operators::KernelInfo>& kernels,
                         std::vector<backend::Program>* programs_out = nullptr) const;

  void registerInvoker(const Invoker& invoker,
                       const operators::NetworkConfig& cfg,
                       const std::string& op,
                       const std::optional<operators::AlgorithmName>& algorithm = std::nullopt) {
    invoker_cache_.registerInvoker({cfg, op}, invoker);
    if (algorithm.has_value()) {
      setNameAndId(cfg, *algorithm, op);
    }
  }

  void setNameAndId(const operators::NetworkConfig& cfg, const operators::AlgorithmName& algorithm, const std::string& op) {
    invoker_cache_.setNameAndId(cfg, algorithm, op);
  }

  std::optional<Invoker> getInvoker(const operators::NetworkConfig& cfg,
                                    const std::optional<operators::Id>& op_id,
                                    const std::optional<operators::AlgorithmName>& algorithm = std::nullopt) const {
    assert(op_id || algorithm);
    if (op_id) {
      return invoker_cache_[std::make_pair(cfg.toString(), op_id->toString())];
    }
    // TODO , algorithm
    return invoker_cache_.getNameId(cfg, *algorithm);
  }

  txdnn::backend::KernelInvoke run(txdnn::backend::Kernel k, bool multi_graph = false) const;
  
  // Program / Code
  backend::Program loadProgram(const fs::path& program_path,
                      std::string params,
                      const std::string& kernel_src,
                      bool force_attach_binary) const;


private:
  std::unique_ptr<HandleImpl> impl_;

  InvokerCache invoker_cache_;
};

} // namespace txdnn

TXDNN_OBJECT_DEFINE(txdnnContext, txdnn::Handle);